import numpy as np
from numpy.lib.function_base import select
from pandas.core.frame import DataFrame
from pandas.core.series import Series
from src.taskSignals import TaskSignals
from src.datahelper import *
from src.backTestingData import BackTestingData
import datetime
from src.strategy.baseStrategy import BaseStrategy
# from onestockinfo import OneStockInfo


class Easystrategy(BaseStrategy):
    def __init__(self, name='easyStrategy') -> None:
        super().__init__()
        self.trans_day_set = set()

    def init_strategy(self, dataHelper: StockDataHelper,
                      backTestingData: BackTestingData, taskSignals: TaskSignals):

        self.taskSignals = taskSignals
        self.data_by_date = dataHelper.get_start_end_dayData(
            backTestingData.data['start_date'], backTestingData.data['end_date'])
        self.trade_day_list = list(self.data_by_date.keys())
        self.position.clear()
        self.portfolio_value = 0.0
        self.trans_period = backTestingData.data['trans_period']
        self.backday = backTestingData.data['back_day']
        self.stocknum = backTestingData.data['stock_nums']
        self.initial_money = backTestingData.data['money']
        self.money = backTestingData.data['money']
        self.service_charge = backTestingData.data['service_charge']
        self.trans_day_set = self.get_trans_period_day_set()  # 所有调仓日的list
        self.daily_return.clear()
        # 初始买入的股票
        self.buy_stocks(
            backTestingData.data['selected_stocks'], self.trade_day_list[0], 'open')
        self.update_portfolio_value(self.trade_day_list[0])

    def get_hold_stocks(self):
        return list(filter(lambda k: self.position[k] != 0, self.position.keys()))

    def get_trans_period_day_set(self):
        """
        :return: 返回所有调仓日list, 隔k个交易日调仓一次
        """
        res = set()
        trans_period = self.trans_period
        for i in range(trans_period - 1, len(self.trade_day_list), trans_period):
            res.add(self.trade_day_list[i])
        return res
    # 调仓日更新

    def update_trans_day_info(self, date):
        stock_num = self.stocknum
        back_day = self.backday

        pre_date = (datetime.datetime.strptime(date, '%Y-%m-%d') +
                    datetime.timedelta(-back_day)).strftime('%Y-%m-%d')
        pre_date = StockDataHelper.find_trade_day(
            pre_date, self.trade_day_list)

        cur_date_series = self.data_by_date[date]['close']
        pre_date_series = self.data_by_date[pre_date]['close']

        # 每支股票的涨幅，其中可能包含NaN的
        rate_series = (cur_date_series - pre_date_series) / pre_date_series
        # 取涨幅最大的k个股票买出， 去除NaN
        df_val = rate_series.dropna().sort_values(
            ascending=False)[:stock_num]

        sell_stocks, buy_stocks = [], []

        # 买入，最多k个， 一定没有NaN
        for stock, val in df_val.items():
            buy_stocks.append(stock)

        # 持有股票全部卖出
        sell_stocks = self.get_hold_stocks()
        # 更新仓位
        self.sell_stocks(sell_stocks, date)
        self.buy_stocks(buy_stocks, date)

    # 交易日的更新

    def update_trade_day_info(self, date):
        pass

    # 运行该策略
    def run(self):
        for i, date in enumerate(self.trade_day_list):
            last = self.portfolio_value + self.money
            if date in self.trans_day_set:
                self.update_trans_day_info(date)
                self.taskSignals.progressBarSignal.emit(
                    int(i / len(self.trade_day_list) * 100))
            else:
                self.update_trade_day_info(date)

            self.update_portfolio_value(date)
            tot = self.portfolio_value + self.money

            # 收益率
            self.daily_return.append((tot - last) / self.initial_money * 100)

            self.log(f'{date}当前总资金{tot}')

        self.log(f'最大回撤{StockDataHelper.max_drawdown(self.daily_return)}')
        self.taskSignals.progressBarSignal.emit(100)

        self.updateBackTestingData()
